# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for kernelized_utils.py."""

import functools

import tensorflow.compat.v2 as tf
from absl.testing import parameterized

from keras.utils import kernelized_utils


def _exact_gaussian(stddev):
    return functools.partial(
        kernelized_utils.exact_gaussian_kernel, stddev=stddev
    )


def _exact_laplacian(stddev):
    return functools.partial(
        kernelized_utils.exact_laplacian_kernel, stddev=stddev
    )


class KernelizedUtilsTest(tf.test.TestCase, parameterized.TestCase):
    @parameterized.named_parameters(
        ("gaussian", _exact_gaussian(stddev=10.0), [[1.0]]),
        ("laplacian", _exact_laplacian(stddev=50.0), [[1.0]]),
    )
    def test_equal_vectors(self, exact_kernel_fn, expected_values):
        """Identical vectors give exactly the identity kernel value."""
        x = tf.constant([0.5, -0.5, -0.5, 0.5])
        y = tf.constant([0.5, -0.5, -0.5, 0.5])
        exact_kernel = exact_kernel_fn(x, y)
        shape = exact_kernel.shape.as_list()
        self.assertLen(shape, 2)
        # x and y are identical and therefore K(x, y) will be precisely equal to
        # the identity value of the kernel.
        self.assertAllClose(expected_values, exact_kernel, atol=1e-6)

    @parameterized.named_parameters(
        ("gaussian", _exact_gaussian(stddev=10.0), [[1.0]]),
        ("laplacian", _exact_laplacian(stddev=50.0), [[1.0]]),
    )
    def test_almost_identical_vectors(self, exact_kernel_fn, expected_values):
        """Almost identical vectors give the identity kernel value."""
        x = tf.constant([1.0, 0.4, -2.1, -1.1])
        y = tf.constant([1.01, 0.39, -2.099, -1.101])
        exact_kernel = exact_kernel_fn(x, y)
        shape = exact_kernel.shape.as_list()
        self.assertLen(shape, 2)
        # x and y are almost identical and therefore K(x, y) will be almost
        # equal to the identity value of the kernel.
        self.assertAllClose(expected_values, exact_kernel, atol=1e-3)

    @parameterized.named_parameters(
        ("gaussian", _exact_gaussian(stddev=1.0), [[0.99], [0.977]]),
        ("laplacian", _exact_laplacian(stddev=5.0), [[0.96], [0.94]]),
    )
    def test_similar_matrices(self, exact_kernel_fn, expected_values):
        """Pairwise "close" vectors give high kernel values (similarity
        scores)."""
        x = tf.constant([1.0, 3.4, -2.1, 0.9, 3.3, -2.0], shape=[2, 3])
        y = tf.constant([1.1, 3.35, -2.05])
        exact_kernel = exact_kernel_fn(x, y)
        shape = exact_kernel.shape.as_list()
        self.assertLen(shape, 2)
        # The 2 rows of x are close to y. The pairwise kernel values (similarity
        # scores) are somewhat close to the identity value of the kernel.
        self.assertAllClose(expected_values, exact_kernel, atol=1e-2)

    @parameterized.named_parameters(
        (
            "gaussian",
            _exact_gaussian(stddev=2.0),
            [[0.997, 0.279], [0.251, 1.0], [0.164, 0.019]],
        ),
        (
            "laplacian",
            _exact_laplacian(stddev=2.0),
            [[0.904, 0.128], [0.116, 1.0], [0.07, 0.027]],
        ),
    )
    def test_matrices_varying_similarity(
        self, exact_kernel_fn, expected_values
    ):
        """Test matrices with row vectors of varying pairwise similarity."""
        x = tf.constant([1.0, 2.0, -2.0, 0.9, 3.3, -1.0], shape=[3, 2])
        y = tf.constant([1.1, 2.1, -2.0, 0.9], shape=[2, 2])
        exact_kernel = exact_kernel_fn(x, y)

        shape = exact_kernel.shape.as_list()
        self.assertLen(shape, 2)
        self.assertAllClose(expected_values, exact_kernel, atol=1e-2)

    @parameterized.named_parameters(
        ("gaussian", _exact_gaussian(stddev=1.0), [[0.0]]),
        ("laplacian", _exact_laplacian(stddev=1.0), [[0.0]]),
    )
    def test_completely_dissimilar_vectors(
        self, exact_kernel_fn, expected_values
    ):
        """Very dissimilar vectors give very low similarity scores."""
        x = tf.constant([1.0, 3.4, -2.1, -5.1])
        y = tf.constant([0.5, 2.1, 1.0, 3.0])
        exact_kernel = exact_kernel_fn(x, y)
        shape = exact_kernel.shape.as_list()
        self.assertLen(shape, 2)
        # x and y are very "far" from each other and so the corresponding kernel
        # value will be very low.
        self.assertAllClose(expected_values, exact_kernel, atol=1e-2)


if __name__ == "__main__":
    tf.test.main()
